import torch
import torch.nn as nn
import torch.nn.functional as F


class DQN(nn.Module):

    def __init__(self, action_dim, device):
        super(DQN, self).__init__() #继承模块，官方规定
        #nn.Conv2d(输入数据的通道数，输出数据的通道数，卷积核大小，步长，零填充)
        self.__conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4, bias=False)   #输入就是当前状态，大小是4
        self.__conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, bias=False)
        self.__conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, bias=False)
        #全连接层输入输出都是二维张量
        #nn.Linear(输入张量大小,输出张量大小)
        self.__fc1 = nn.Linear(64*7*7, 512)
        self.__fc2 = nn.Linear(512, action_dim) #输出就是动作
        self.__device = device

    def forward(self, x):
        x = x / 255.#x是一个状态，大小为4
        x = F.relu(self.__conv1(x))  #激励函数 relu 负值取零，保留正值，处理神经网络输出的结果
        x = F.relu(self.__conv2(x))
        x = F.relu(self.__conv3(x))
        x = F.relu(self.__fc1(x.view(x.size(0), -1))) #x.size(0)是batch size，view函数相当于numpy的reshape（-1指自适应）
        return self.__fc2(x)

    @staticmethod
    def init_weights(module):
        if isinstance(module, nn.Linear):
            torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
            module.bias.data.fill_(0.0)
        elif isinstance(module, nn.Conv2d):
            torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")

class DuelingDQN(nn.Module):

    def __init__(self, action_dim, device):
        super(DuelingDQN, self).__init__() #继承模块，官方规定
        self.action_dim = action_dim
        #nn.Conv2d(输入数据的通道数，输出数据的通道数，卷积核大小，步长，零填充)
        self.__conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4, bias=False)   #输入就是当前状态，大小是4
        self.__conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, bias=False)
        self.__conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, bias=False)
        #全连接层输入输出都是二维张量
        #nn.Linear(输入张量大小,输出张量大小)
        #self.__fc1 = nn.Linear(64*7*7, 512)
        #self.__fc2 = nn.Linear(512, action_dim) #输出就是动作
        
        self.__fc1_adv = nn.Linear(7*7*64, 512)
        self.__fc1_val = nn.Linear(7*7*64, 512)

        self.__fc2_adv = nn.Linear(512, action_dim)
        self.__fc2_val = nn.Linear(512, 1)
        
        self.__device = device

    def forward(self, x):
        x = x / 255.#x是一个状态，大小为4
        x = F.relu(self.__conv1(x))  #激励函数 relu 负值取零，保留正值，处理神经网络输出的结果
        x = F.relu(self.__conv2(x))
        x = F.relu(self.__conv3(x))
        x = x.view(x.size(0), -1)
        
        adv = F.relu(self.__fc1_adv(x))
        val = F.relu(self.__fc1_val(x))
        
        adv = self.__fc2_adv(adv)
        val = self.__fc2_val(val).expand(x.size(0), self.action_dim)
        
        x = val + adv - adv.mean(1).unsqueeze(1).expand(x.size(0), self.action_dim)
        return x

    @staticmethod
    def init_weights(module):
        if isinstance(module, nn.Linear):
            torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
            module.bias.data.fill_(0.0)
        elif isinstance(module, nn.Conv2d):
            torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")